package com.example.mcpclient;

import com.example.mcpclient.service.ShopService;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
import lombok.extern.slf4j.Slf4j;
import org.apache.http.HttpHost;
import org.elasticsearch.client.RestClient;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.tool.method.MethodToolCallbackProvider;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStore;
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStoreOptions;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.boot.autoconfigure.elasticsearch.ElasticsearchProperties;
import org.springframework.context.annotation.Bean;

import java.net.URI;
import java.net.URISyntaxException;

import static org.springframework.ai.vectorstore.elasticsearch.SimilarityFunction.cosine;

@SpringBootApplication
@Slf4j
public class McpClientApplication {

    public static void main(String[] args) {
        SpringApplication.run(McpClientApplication.class, args);
        log.info("====================== mcp-client ======================");
    }


    /**
     * 对话客户端
     * <p>
     * 有打印了工具回调的方法名称
     * </p>
     */
    @Bean
    public ChatClient chatClient(
            ChatClient.Builder chatClientBuilder,
            ToolCallbackProvider tools,
            ShopService shopService,
            Advisor groupIdAdvisor,
            Advisor localQuestionAdvisor,
            Advisor esQuestionAdvisor,
            Advisor promptChatMemoryAdvisor,
            Advisor simpleLoggerAdvisor) {

        // 查看工具回调信息
        ToolCallback[] callbacks = tools.getToolCallbacks();
        log.info("Tool callbacks count: {}", callbacks.length);
        for (ToolCallback callback : callbacks) {
            log.info("Tool: {} - {}", callback.getToolDefinition().name(), callback.getToolDefinition().name());
        }

        return chatClientBuilder
                .defaultTools(shopService)
                .defaultAdvisors(groupIdAdvisor,
                        localQuestionAdvisor,
                        esQuestionAdvisor,
                        promptChatMemoryAdvisor,
                        simpleLoggerAdvisor)
                .build();
    }

    /**
     * 本机内存向量库
     */
    @Bean
    public VectorStore localVectorStore(EmbeddingModel embeddingModel) {
        return SimpleVectorStore.builder(embeddingModel)
                .build();
    }

    /**
     * rag向量增强 _ 内存向量库
     */
    @Bean
    public Advisor localQuestionAdvisor(VectorStore localVectorStore) {
        return QuestionAnswerAdvisor.builder(localVectorStore)
                .searchRequest(SearchRequest.builder()
                        .similarityThreshold(0.1d).topK(5).build())
                .build();
    }


    /**
     * es向量库
     */
    @Bean
    public VectorStore esVectorStore(RestClient restClient, EmbeddingModel embeddingModel) {
        ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
        // defaults to "spring-ai-document-index"
        options.setIndexName("custom-index");
        // Optional: defaults to COSINE
        options.setSimilarity(cosine);
        // Optional: defaults to model dimensions or 1536
        options.setDimensions(1536);

        return ElasticsearchVectorStore.builder(restClient, embeddingModel)
                .options(options)
                .initializeSchema(true)
                .batchingStrategy(new TokenCountBatchingStrategy())
                .build();
    }

    /**
     * es客户端
     */
    @Bean
    public RestClient restClient(ElasticsearchProperties properties) {
        String uriString = properties.getUris().get(0);
        try {
            URI uri = new URI(uriString);
            return RestClient.builder(new HttpHost(uri.getHost(), uri.getPort(), uri.getScheme()))
                    .build();
        } catch (URISyntaxException e) {
            throw new IllegalArgumentException("Invalid Elasticsearch URI: " + uriString, e);
        }
    }

    /**
     * rag向量增强 - es
     */
    @Bean
    public Advisor esQuestionAdvisor(VectorStore esVectorStore) {
        return QuestionAnswerAdvisor.builder(esVectorStore)
                .searchRequest(SearchRequest.builder()
                        .similarityThreshold(0.1d).topK(5).build())
                .build();
    }

    /**
     * 指定tool方法, 类似于mcp的tool
     */
    @Bean
    public ToolCallbackProvider tools(ShopService shopService) {
        return MethodToolCallbackProvider.builder()
                .toolObjects(shopService)
                .build();
    }


    /**
     * 打印日志用json
     */
    @Bean
    public Advisor simpleLoggerAdvisor() {
        return new SimpleLoggerAdvisor(
                request -> {
                    try {
                        ObjectMapper objectMapper = new ObjectMapper();
                        // 注册JavaTimeModule来处理Java 8日期时间类型
                        objectMapper.registerModule(new JavaTimeModule());
                        // 禁用将时间序列化为时间戳，使其更可读
                        objectMapper.disable(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS);
                        return objectMapper.writeValueAsString(request);
                    } catch (Exception e) {
                        log.error("Failed to serialize request to JSON", e);
                        return "入参打印日志异常";
                    }
                },
                response -> {
                    try {
                        ObjectMapper objectMapper = new ObjectMapper();
                        // 注册JavaTimeModule来处理Java 8日期时间类型
                        objectMapper.registerModule(new JavaTimeModule());
                        // 禁用将时间序列化为时间戳，使其更可读
                        objectMapper.disable(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS);
                        return objectMapper.writeValueAsString(response);
                    } catch (Exception e) {
                        log.error("Failed to serialize response to JSON", e);
                        return "返回值打印日志异常";
                    }
                }, 0);
    }

    /**
     * 对话历史记忆
     */
    @Bean
    public Advisor promptChatMemoryAdvisor(ChatMemory chatMemory) {
        return PromptChatMemoryAdvisor.builder(chatMemory).build();
    }

    /**
     * rag向量增强 更复杂的场景
     */
//    @Bean
//    public Advisor retrievalAugmentationAdvisor(SimpleVectorStore vectorStore) {
//        return RetrievalAugmentationAdvisor.builder()
//                .documentRetriever(VectorStoreDocumentRetriever.builder()
//                        .similarityThreshold(0.50)
//                        .vectorStore(vectorStore)
//                        .build())
//                .build();
//    }


}
